/*
 * Copyright 2017 Baiyyy
 * @(#)DataRulePaginationInterceptor.java
 */

package com.byyy.ccts.modules.common.interceptor;

import com.baomidou.mybatisplus.toolkit.PluginUtils;
import com.byyy.ccts.core.security.shiro.authz.annotation.DataPermission;
import com.byyy.ccts.core.utils.StringUtils;
import com.byyy.ccts.modules.sys.entity.Cangku;
import com.byyy.ccts.modules.sys.entity.User;
import com.byyy.ccts.modules.sys.utils.UserUtils;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ResultMap;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.List;
import java.util.Properties;

/**
 * 自定义数据权限控制插件
 */
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class DataPermissionInterceptor implements Interceptor {

    private static String SELECT = "SELECT";
    private static String FROM = "FROM";
    private static String PERMISSION_CANGKU = "CANGKU";

    /**
     * 数据权限控制插件
     *
     * 针对于返回类型，解析类字段是否含有@DataPermission注解，存在则解析权限拼接SQL
     * @param invocation
     * @return
     * @throws Throwable
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaStatementHandler = SystemMetaObject.forObject(statementHandler);
        // 先判断是不是SELECT操作
        MappedStatement mappedStatement = (MappedStatement) metaStatementHandler.getValue("delegate.mappedStatement");
        if (!SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType())) {
            return invocation.proceed();
        }
        // 针对定义了rowBounds，做为mapper接口方法的参数
        BoundSql boundSql = (BoundSql) metaStatementHandler.getValue("delegate.boundSql");
        String originalSql = boundSql.getSql();

        // 定义变量
        List<ResultMap> resultMaps = mappedStatement.getResultMaps();
        Class resultClz = null;
        Field[] resultFld = null;
        String permissionName = null;
        Class<DataPermission> dataPermissionClz = DataPermission.class;
        StringBuilder dataPermissionSqlSb = new StringBuilder();
        if (resultMaps == null || resultMaps.size() < 1) {
            return invocation.proceed();
        }

        // 获取返回对象，判断返回对象字段是否存在数据权限注解
        for (ResultMap resultMap : resultMaps) {
            resultClz = resultMap.getType();
            resultFld = resultClz.getDeclaredFields();
            for (Field field : resultFld) {
                if (field.isAnnotationPresent(dataPermissionClz)) {
                    permissionName = field.getAnnotation(DataPermission.class).value();
                    permissionName = StringUtils.isNotEmpty(permissionName) ? permissionName : field.getName();

                    // 判断SQL返回值是否存在该字段，如果不存在，就不计算入数据权限内
                    if(!containsPermissionInResult(originalSql, permissionName)) {
                        continue;
                    }

                    // 获取用户角色，如果不存在用户角色，则不需要数据权限验证
                    User user = UserUtils.getUser();
                    if (StringUtils.isEmpty(user.getUsername())) {
                        return invocation.proceed();
                    }

                    // 获取用户角色，判断用户角色中对应的数据权限
                    if (PERMISSION_CANGKU.equals(permissionName.toUpperCase())) {
                        dataPermissionSqlSb.append(" AND ").append(generateDataPermissionSQL());
                    }
                }
            }
        }

        // 不存在数据权限继续执行拦截，存在数据权限则封装SQL之后继续执行拦截
        if (StringUtils.isEmpty(dataPermissionSqlSb.toString())) {
            return invocation.proceed();
        }

        originalSql = "SELECT T.* FROM (" + originalSql + ") T WHERE 1 = 1" + dataPermissionSqlSb.toString();
        metaStatementHandler.setValue("delegate.boundSql.sql", originalSql);

        // 调用父类的拦截方法
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(Properties prop) {
    }

    /**
     * 判断SQL返回结果中是否包含数据权限字段
     *
     * @param sql
     * @param permissionName
     * @return
     */
    private boolean containsPermissionInResult(String sql, String permissionName) {
        StringBuilder sb = new StringBuilder(sql.toUpperCase());
        int startIndex = 0;
        int fromIndex = 0;
        int nextIndex = 0;

        // 循环遍历SELECT至FROM之间的返回值，判断是否包含数据权限字段
        while (sb.indexOf(SELECT, nextIndex) >= 0) {
            startIndex = sb.indexOf(SELECT, nextIndex);
            fromIndex = sb.indexOf(FROM, nextIndex);
            nextIndex = fromIndex + 4;
            // 遍历结束
            if (startIndex < 0 || fromIndex < 0) {
                return false;
            }
            // 判断数据权限是否在返回值中
            if (sb.substring(startIndex, fromIndex).contains(permissionName.toUpperCase())) {
                return true;
            }
        }

        return false;
    }

    /**
     * 根据数据权限名称，生成SQL
     * @return
     */
    private String generateDataPermissionSQL() {
        StringBuilder sql = new StringBuilder();
        sql.append(PERMISSION_CANGKU).append(" IN (");

        // 获取用户角色对应的数据权限
        List<Cangku> cangkuList = UserUtils.getUserCangkuList();
        if (cangkuList == null || cangkuList.size() < 1) {
            sql.append("'')");
            return sql.toString();
        }

        // 生成对应的SQL
        for (Cangku cangku : cangkuList) {
            sql.append("'").append(cangku.getCangku().toUpperCase()).append("', ");
        }
        sql.deleteCharAt(sql.lastIndexOf(",")).append(")");

        return sql.toString();
    }

}
